package cn.com.zetatech.loader.common.db;

import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.session.ExecutorType;
import org.apache.ibatis.session.SqlSession;
import org.apache.ibatis.session.SqlSessionFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.jdbc.core.BatchPreparedStatementSetter;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.datasource.DataSourceUtils;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BiConsumer;

/**
 * @author zhengbs
 */
@Component
@Slf4j
public class DbService {

    @Autowired
    private JdbcTemplate jdbcTemplate;
    @Autowired
    private SqlSessionFactory sqlSessionFactory;


    public String getDatabaseId() {
        DataSource dataSource = jdbcTemplate.getDataSource();
        Connection connection = null;
        try {
            connection = DataSourceUtils.getConnection(dataSource);
            String databaseProductName = connection.getMetaData().getDatabaseProductName();
            DataSourceUtils.releaseConnection(connection, dataSource); // 手动关闭数据库连接并返回到连接池
            return databaseProductName;
        } catch (Exception e) {
            log.error(e.getMessage(), e);
            if (connection != null) {
                DataSourceUtils.releaseConnection(connection, dataSource); // 手动关闭数据库连接并返回到连接池
            }
        }
        return null;
    }

    public <T> void batchInsert(String sql, List<T> entities, int batchSize, int threadCount, BiConsumer<T, PreparedStatement> setter) {
        if(entities.size()>0) {
            AtomicReference<String> message = new AtomicReference<>("");
            if(entities.size() <= batchSize) {
                batchInsert(sql, entities, setter);
            } else {
                ExecutorService executor = Executors.newFixedThreadPool(threadCount);

                CountDownLatch latch = new CountDownLatch(threadCount);

                for (int i = 0; i < entities.size(); i += batchSize) {
                    List<T> subData = entities.subList(i, Math.min(i + batchSize, entities.size()));

                    executor.submit(() -> {
                        long start_time = System.currentTimeMillis();
                        DataSource dataSource = jdbcTemplate.getDataSource();
                        Connection connection = null;
                        try {
                            connection = DataSourceUtils.getConnection(dataSource);
                            log.info("获取{}。耗时{}ms", connection, System.currentTimeMillis() - start_time);
                            start_time = System.currentTimeMillis();
                            connection.setAutoCommit(false); // 关闭自动提交
                            jdbcTemplate.batchUpdate(sql, new BatchPreparedStatementSetter() {
                                @Override
                                public void setValues(PreparedStatement ps, int i) throws SQLException {
                                    T entity = subData.get(i);
                                    setter.accept(entity, ps); // 调用setter方法设置参数值
                                }

                                @Override
                                public int getBatchSize() {
                                    return subData.size();
                                }
                            });
                            connection.commit(); // 提交事务
                            log.info("{}笔数据，已提交{}，耗时{}ms", subData.size(), subData.size(), System.currentTimeMillis() - start_time);
                            DataSourceUtils.releaseConnection(connection, dataSource); // 手动关闭数据库连接并返回到连接池
                        } catch (Exception e) {
                            log.error(e.getMessage(), e);
                            if (connection != null) {
                                try {
                                    connection.rollback(); // 发生异常，回滚事务
                                } catch (SQLException ex) {
                                    log.error(e.getMessage(), e);
                                }
                                DataSourceUtils.releaseConnection(connection, dataSource); // 手动关闭数据库连接并返回到连接池
                            }
                        } finally {
                            latch.countDown();
                        }
                    });
                }
                try {
                    latch.await();
                    executor.shutdown();
                } catch (InterruptedException e) {
                    log.error(e.getMessage(), e);
                }

            }
            if (!StringUtils.isEmpty(message.get())) {
                throw new RuntimeException(message.get());
            }
        }
    }

    public <T> void batchInsert(String sql, List<T> entities, BiConsumer<T, PreparedStatement> setter) {
        if(entities.size()>0) {
            long start_time = System.currentTimeMillis();
            DataSource dataSource = jdbcTemplate.getDataSource();
            Connection connection = null;
            try {
                connection = DataSourceUtils.getConnection(dataSource);
                connection.setAutoCommit(false); // 关闭自动提交
                jdbcTemplate.batchUpdate(sql, new BatchPreparedStatementSetter() {
                    @Override
                    public void setValues(PreparedStatement ps, int i) throws SQLException {
                        T entity = entities.get(i);
                        setter.accept(entity, ps); // 调用setter方法设置参数值
                    }

                    @Override
                    public int getBatchSize() {
                        return entities.size();
                    }
                });
                connection.commit(); // 提交事务
                log.info("{}笔数据，已提交{}，耗时{}ms", entities.size(), entities.size(), System.currentTimeMillis() - start_time);
                DataSourceUtils.releaseConnection(connection, dataSource); // 手动关闭数据库连接并返回到连接池
            } catch (Exception e) {
                log.error(e.getMessage(), e);
                if (connection != null) {
                    try {
                        connection.rollback(); // 发生异常，回滚事务
                    } catch (SQLException ex) {
                        log.error(e.getMessage(), e);
                    }
                    DataSourceUtils.releaseConnection(connection, dataSource); // 手动关闭数据库连接并返回到连接池
                }
                throw new RuntimeException(e);
            }
        }
    }

    public <T> void batchInsert(List<T> entities, int batchSize, int threadCount, BiConsumer<SqlSession, T> consumer) {
        if(entities.size()>0) {
            AtomicReference<String> message = new AtomicReference<>("");
            if(entities.size() <= batchSize) {
                batchInsert(entities, consumer);
            } else {
                ExecutorService executor = Executors.newFixedThreadPool(threadCount);

                CountDownLatch latch = new CountDownLatch(threadCount);

                for (int i = 0; i < entities.size(); i += batchSize) {
                    List<T> subData = entities.subList(i, Math.min(i + batchSize, entities.size()));
                    executor.submit(() -> {
                        AtomicLong count = new AtomicLong();
                        long start_time = System.currentTimeMillis();
                        SqlSession sqlSession = sqlSessionFactory.openSession(ExecutorType.BATCH);
                        log.info("获取{}。耗时{}ms", sqlSession, System.currentTimeMillis() - start_time);
                        try {
                            start_time = System.currentTimeMillis();
                            int size = subData.size();
                            int j = 1;
                            for (T element : subData) {
                                consumer.accept(sqlSession, element);
                                count.addAndGet(1);
                                if ((j % batchSize == 0) || j == size) {
                                    sqlSession.flushStatements();
                                    log.info("{}笔数据，已提交{}，耗时{}ms", size, count.get(), System.currentTimeMillis() - start_time);
                                    start_time = System.currentTimeMillis();
                                }
                                j++;
                            }
                        } catch (Exception e) {
                            message.set(e.getMessage());
                        } finally {
                            latch.countDown();
                            sqlSession.commit(true);
                            sqlSession.clearCache();
                            sqlSession.close();
                        }
                    });
                }
                try {
                    latch.await();
                    executor.shutdown();
                } catch (InterruptedException e) {
                    log.error(e.getMessage(), e);
                }

            }
            if (!StringUtils.isEmpty(message.get())) {
                throw new RuntimeException(message.get());
            }
        }
    }


    public <T> void batchInsert(List<T> entities, BiConsumer<SqlSession, T> consumer) {
        if(entities.size()>0) {
            String errorMsg = null;
            long start_time = System.currentTimeMillis();
            SqlSession sqlSession = sqlSessionFactory.openSession(ExecutorType.BATCH);
            try {
                for (T element : entities) {
                    consumer.accept(sqlSession, element);
                }
                sqlSession.flushStatements();
                log.info("{}笔数据，已提交{}，耗时{}ms", entities.size(), entities.size(), System.currentTimeMillis() - start_time);
            } catch (Exception e) {
                errorMsg = e.getMessage();
            } finally {
                sqlSession.commit(true);
                sqlSession.clearCache();
                sqlSession.close();
            }
            if (!StringUtils.isEmpty(errorMsg)) {
                throw new RuntimeException(errorMsg);
            }
        }
    }

    public <T> void batchInsertWithRetry(List<T> entities, BiConsumer<SqlSession, T> consumer, int maxRetries, long retryInterval) {
        if (entities.isEmpty()) return;

        String errorMsg = null;
        long start_time = System.currentTimeMillis();
        int retryAttempts = 0;
        SqlSession sqlSession = null;

        try {
            while (retryAttempts <= maxRetries) {
                retryAttempts++;

                sqlSession = sqlSessionFactory.openSession(ExecutorType.BATCH);
                for (T element : entities) {
                    consumer.accept(sqlSession, element);
                }
                sqlSession.flushStatements();
                sqlSession.commit(true);
                log.info("{}笔数据，已提交{}，耗时{}ms", entities.size(), entities.size(), System.currentTimeMillis() - start_time);
                return; // 成功后直接返回
            }
        } catch (Exception e) {
            errorMsg = e.getMessage();
            log.error("批量插入第{}次尝试失败，原因: {}", retryAttempts, errorMsg);
            if (sqlSession != null) {
                sqlSession.rollback(true);
            }
            if (retryAttempts <= maxRetries) { // 这里应确保在最后一次尝试后再休眠
                try {
                    TimeUnit.MILLISECONDS.sleep(retryInterval);
                } catch (InterruptedException ie) {
                    Thread.currentThread().interrupt();
                    throw new RuntimeException("插入数据重试过程中被中断", ie);
                }
            }
        } finally {
            if (sqlSession != null) {
                try {
                    sqlSession.clearCache();
                } finally {
                    sqlSession.close();
                }
            }
        }
        log.info("批量插入数据失败，经过" + maxRetries + "次重试后依旧失败，错误信息: " + errorMsg);
//        throw new RuntimeException("批量插入数据失败，经过" + maxRetries + "次重试后依旧失败，错误信息: " + errorMsg);
    }

}
